package invertedindex;

import java.io.*;
import java.sql.*;
import java.util.*;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

public class Index {
	static Connection conn;
    static PreparedStatement stmt_insert_into_postings; 
    static PreparedStatement stmt_select_all_from_postings;
    static PreparedStatement stmt_insert_into_docmag;
    static PreparedStatement stmt_select_df_from_postings;
    static PreparedStatement stmt_select_docinfo_from_docmag;
    static PreparedStatement stmt_select_docinfos_from_docmag;
    static PreparedStatement stmt_insert_into_documents;
    static PreparedStatement stmt_select_document_from_documents;

    static int batch_size = 10;
    static boolean connected;
    
    static {
    	try { 
    		conn = ConnectionManager.getConnection();
    		connected = true;
    		
    		String postings_table = "postings";
    		stmt_insert_into_postings = conn.prepareStatement("INSERT INTO "+postings_table+"(word, df, docidfreqbytestream) VALUES(?,?,?)");
    		stmt_select_all_from_postings = conn.prepareStatement("SELECT word, df, docidfreqbytestream FROM "+postings_table);
    		stmt_select_df_from_postings = conn.prepareStatement("SELECT df FROM "+postings_table+" WHERE word=?");

    		String docmag_table = "docmag_clustered";
    		// insert a row into docmag table
    		stmt_insert_into_docmag = conn.prepareStatement("INSERT INTO "+docmag_table+" (docid, vectormagnitude, maxf) VALUES(?,?,?)");
    		// select single row from docmag table
    		stmt_select_docinfo_from_docmag = conn.prepareStatement("SELECT vectormagnitude, maxf FROM "+docmag_table+" WHERE docid=?");
    		
    		// Select [batch_size] rows from docmag at a time
    		stmt_select_docinfos_from_docmag = conn.prepareStatement("SELECT docid, vectormagnitude, maxf FROM "+docmag_table+" WHERE docid IN " + makeQueryPlaceholders(batch_size));
    		
    		// Insert into documents table
    		stmt_insert_into_documents = conn.prepareStatement("INSERT INTO documents(docid, doctitle, doctext) VALUES(?,?,?)");

    		// Select title and text from documents table
    		stmt_select_document_from_documents = conn.prepareStatement("SELECT doctitle, doctext FROM documents WHERE docid = ?");
    	}
    	catch (Exception e) {
    		System.err.println(e);
    		connected = false;
    	}; 		
    }

    /**
     * Forces user to type "y" to a prompt before overwriting database contents.
     * Returns true if user responded with "y".
     */
    public static boolean askBeforeWrite(String prompt) throws IOException {
	    System.err.println(prompt);
		BufferedReader stdin = new BufferedReader(new InputStreamReader(System.in));
	    String response = stdin.readLine();
	    if(response.equals("y")) {
	    	return true;
	    }
	    System.out.println("Quitting...");
	    return false;
    }
    
    /**
     * Fills the postings table, given the sorted list of word -> doc_id pairs.
     * @param sortedWordDocListFilename
     * @throws Exception
     */
	public static void fillWordTable (String sortedWordDocListFilename) throws Exception {
		if(!askBeforeWrite("Are you sure? This will modify the postings table. [y/N]")) {
			return;
		}
		
		StopWatch sw = new StopWatch("Filling the postings table");

		// open file
        BufferedReader in = new BufferedReader(new FileReader(sortedWordDocListFilename) );

		// create docid_freq_map as {doc_id -> frequency of single word occurrence}
        SortedMap<Integer, Integer> docid_freq_map = new TreeMap<Integer, Integer>();

		// for each line in file:
        String line;
        String lastWord = null;
        while ( (line = in.readLine()) != null )  {
        	String[] elements = line.split(" ");
        	String word = elements[0];
        	int doc_id = Integer.parseInt(elements[1]);
    		
        	// If word is the same or first word, update docid_freq_map
        	if(word.equals(lastWord) || lastWord == null) {

        	} else {
        		System.out.println(lastWord);
        		
           		// New word found - insert data and clear docid_freq_map
    			stmt_insert_into_postings.setString(1, lastWord);
    			stmt_insert_into_postings.setInt(2, docid_freq_map.size());
    			stmt_insert_into_postings.setBytes(3, VB.VBENCODE( Delta.sortedmap2deltalist(docid_freq_map) ) ); 					
    			stmt_insert_into_postings.execute();
    			
        		docid_freq_map.clear();
        	}
        	
        	// Increment or create counter for occurrences of this word in this document
       		if(!docid_freq_map.containsKey(doc_id)) {
    			docid_freq_map.put(doc_id, 1);
    		} else {
    			docid_freq_map.put(doc_id, docid_freq_map.get(doc_id) + 1);
    		}

       		lastWord = word;
        }
        in.close();

        conn.commit();
        ConnectionManager.returnConnection(conn);
	
        sw.stop();
    }    
    
    /**
     * Creates the sorted word->document frequency text file from the postings table.
     * @param DocWordFreqDfListFilename the filename to write to.
     * @throws Exception
     */
	public static void fillDocWordFreqDfList (String DocWordFreqDfListFilename) throws Exception {
		StopWatch sw = new StopWatch("doc_word_freq_df.txt");

		//Creating the DocWordFreqDfList output file
    	BufferedWriter out = new BufferedWriter( new FileWriter(DocWordFreqDfListFilename) );
        
    	ResultSet rset = stmt_select_all_from_postings.executeQuery();
    	while (rset.next()) {
    		String word = rset.getString(1);
    		System.out.println(word);
    		Integer df = rset.getInt(2);
    		byte[] bytestream = rset.getBytes(3);
    		SortedMap<Integer, Integer> docid_freq_map = Delta.deltalist2sortedmap( VB.VBDECODE(bytestream) );
    		for (Integer docid : docid_freq_map.keySet()) {
    			out.write(docid + " " + word + " " + docid_freq_map.get(docid) + " " + df + "\n");
    		}
    	}
        
    	out.close();
        ConnectionManager.returnConnection(conn);
        sw.stop();
	}
	
	
    /**
     * Fills the docmag table given the sorted word -> document frequency text file.
     * @param sortedDocWordDfListFilename
     * @param numDocs The total number of documents
     * @throws Exception
     */
	public static void fillDocMagTable (String sortedDocWordDfListFilename, int numDocs) throws Exception {
		if(!askBeforeWrite("Are you sure? This will modify the docmag table. [y/N]")) {
			return;
		}

		StopWatch sw = new StopWatch("Filling docmag table");
		
		// Open file
        BufferedReader in = new BufferedReader( new FileReader(sortedDocWordDfListFilename) );
        
        Map<String, Integer> word_freq_map = new HashMap<String, Integer>();
        Map<String, Integer> word_df_map = new HashMap<String, Integer>();

        String line;
        int last_doc_id = -1;
        
        while ( (line = in.readLine()) != null )  {
        	String[] elements = line.split(" ");
        	int doc_id = Integer.parseInt(elements[0]);
        	String word = elements[1];
        	int freq = Integer.parseInt(elements[2]);
        	int df = Integer.parseInt(elements[3]);
        	
        	if(doc_id != last_doc_id && last_doc_id != -1) {
        		// New document, insert into docmag table and clear freq_map
    			stmt_insert_into_docmag.setInt(1, doc_id);
    			stmt_insert_into_docmag.setDouble(2, Index.magnitude(word_freq_map, word_df_map, numDocs));
    			stmt_insert_into_docmag.setInt(3, Index.maxf(word_freq_map));
    			stmt_insert_into_docmag.execute();
    			
        		word_freq_map.clear();
        		System.out.println(doc_id + " / " + numDocs);
        	}
        	
        	// Update freq_map and df_map
       		word_freq_map.put(word, freq);
       		if(!word_df_map.containsKey(word)) {
       			word_df_map.put(word, df);
       		}
       		
       		last_doc_id = doc_id;
        }
        	
        conn.commit();
        ConnectionManager.returnConnection(conn);
        sw.stop();
    }
	
	public static int maxf (Map<String, Integer> word_freq_map) {
		int maxf_res = 0;
		for (String word : word_freq_map.keySet()) {
			if (maxf_res < word_freq_map.get(word))
				maxf_res = word_freq_map.get(word);
		}
		return maxf_res;
	}	
	
	public static double tfidf (int f, int maxf, int df, int numDocs) {
    	if(f == 0 || df == 0 || maxf == 0) {
    		return 0;
    	}

		double tf = f/(double)maxf;
    	double idf = Math.log(numDocs/(double)df);
    	return tf*idf; 		
	}
	
	public static double magnitude (Map<String, Integer> word_freq_map, int numDocs) throws Exception 
	{
		Map<String, Integer> word_df_map = new HashMap<String, Integer>();
		for (String word : word_freq_map.keySet()) 
			word_df_map.put( word, Index.retrieveDf(word) );
		
		return Index.magnitude(word_freq_map, word_df_map, numDocs);	
	}
	
	public static double magnitude (Map<String, Integer> word_freq_map,
									Map<String, Integer> word_df_map,
									int numDocs) throws Exception {		
		double magsquare = 0;
		int maxf = maxf(word_freq_map);
		
		for(String word : word_freq_map.keySet()) {
			int f = word_freq_map.get(word);
			int df = word_df_map.get(word);
			double tfidf = tfidf(f, maxf, df, numDocs);
			magsquare += Math.pow(tfidf, 2);
		}		
		
		return Math.sqrt(magsquare); 
	}

	public static int retrieveDf (String word) throws Exception {
		if(!connected) {
			throw new Exception("Failed to connect to database");
		}

		int df = 0;
		stmt_select_df_from_postings.setString(1, word);
    	ResultSet rset = stmt_select_df_from_postings.executeQuery();
    	if (rset.next()) { 
    		df = rset.getInt(1); 
    	}
		return df;
	}
	
	public static DocInfo retrieveDocInfo (Integer docid) throws Exception {
		if(!connected) {
			throw new Exception("Failed to connect to database");
		}

		DocInfo docinfo = new DocInfo();

	    stmt_select_docinfo_from_docmag.setInt(1, docid);
	    
    	ResultSet rset = stmt_select_docinfo_from_docmag.executeQuery();
    	if (rset.next()) { 
    		docinfo.docid = docid;
    		docinfo.magnitude = rset.getDouble(1); 
    		docinfo.maxf = rset.getInt(2);
    	} else {
    		throw new Exception("doc_id " + docid + " not found in docmag table!");
    	}

		return docinfo;
	}
	
	/**
	 * Retrieves a List of DocInfo instances given a Set of Integer document IDs. 
	 * Uses query batching to improve performance over using one query per row.
	 * @param docid_set A set of the document IDs to retrieve from the docmag table.
	 * @return A List of DocInfo instances representing the tuples from the docmag table.
	 * @throws Exception
	 */
	public static List<DocInfo> retrieveDocInfoBatch(Set<Integer> docid_set) throws Exception {
		if(!connected) {
			throw new Exception("Failed to connect to database");
		}
		List<DocInfo> result = new LinkedList<DocInfo>();
		Integer[] docids = docid_set.toArray(new Integer[0]);
		int retrieveCount = 0;
		
		int num_batches = docids.length / batch_size;
		int docid = -1;
		
		// Retrieve rows from docmag [batch_size] rows at a time
		for(int batch=0; batch<num_batches; batch++) {
			for(int param=0; param<batch_size; param++) {
				docid = docids[batch * batch_size + param];
				stmt_select_docinfos_from_docmag.setInt(param + 1, docid);
			}

			ResultSet rset = stmt_select_docinfos_from_docmag.executeQuery();
			while (rset.next()) { 
				DocInfo docinfo = new DocInfo();
				docinfo.docid = rset.getInt(1);
				docinfo.magnitude = rset.getDouble(2); 
				docinfo.maxf = rset.getInt(3);
				result.add(docinfo);
				retrieveCount++;
			}
		}
		
		// Retrieve remaining docmag rows one at a time
		for(int i=num_batches*batch_size; i<docids.length; i++) {
			docid = docids[i];
			result.add(retrieveDocInfo(docid));
			retrieveCount++;
		}
		
		if(retrieveCount != docids.length) {
			throw new Exception("Missing document ID's: " + (docids.length - retrieveCount));
		}
		return result;
	}
	
	public static Document retrieveDocument (Integer docid) throws Exception {
		if(!connected) {
			throw new Exception("Failed to connect to database");
		}

		Document document;
	    stmt_select_document_from_documents.setInt(1, docid);
	    
    	ResultSet rset = stmt_select_document_from_documents.executeQuery();
    	if (rset.next()) { 
    		document = new Document(rset.getString(1), rset.getString(2));
    	} else {
    		throw new Exception("doc_id " + docid + " not found in document table!");
    	}

		return document;
	}

	/**
	 * Utility method for making batched queries. 
	 * e.g. SELECT * FROM docmag WHERE id in (?, ?, ?, ?)
	 * @param size the number of placeholders used in the query
	 * @return
	 */
	public static String makeQueryPlaceholders(int size) {
		if(size <= 0) { return "()"; }
		
		StringBuilder b = new StringBuilder("(");
		for(int i=0; i<size-1; i++) {
			b.append("?,");
		}
		b.append("?)");
		return b.toString();
	}

	public static void fillDocumentsTable(String documentsFolder) throws Exception {
		StopWatch sw = new StopWatch("Filling documents table");
        File dir = new File(documentsFolder);
        File[] files = dir.listFiles(); 
        String newline = System.getProperty("line.separator");
        // Match header
        Pattern titlePattern = Pattern.compile("<H2>([^<]+)</H2>");
        
    	for (File file : files) {
            System.out.println("Processing file: " + file.getName());
            Integer docid = Integer.parseInt(file.getName().replaceAll(".html", ""));

            BufferedReader in = new BufferedReader( new FileReader(file) );
            String line;
            StringBuilder doctext = new StringBuilder();
            
            while ( (line = in.readLine()) != null )  {
            	doctext.append(line);
            	doctext.append(newline);
            }
            
            String doctitle;
            Matcher titleMatcher = titlePattern.matcher(doctext);
            // skip over first match (the date)
            titleMatcher.find();
            // extract the title from the second header element
            if(titleMatcher.find()) {
            	doctitle = titleMatcher.group(1);
            } else {
            	doctitle = "";
            }
            
			stmt_insert_into_documents.setInt(1, docid);
			stmt_insert_into_documents.setString(2, doctitle);
			stmt_insert_into_documents.setString(3, doctext.toString());

			stmt_insert_into_documents.execute();

            in.close();
        }
    	
        conn.commit();
        ConnectionManager.returnConnection(conn);

    	sw.stop();
        System.out.println("Processed " + files.length + " HTML files");
	}
	
    public static void main(String[] args) throws Exception {
        //Call the following only once each in the following oder.

    	//DO: sort -k1,1 -k2n,2 word_doc_list.txt > sorted_word_doc_list.txt
    	//Index.fillWordTable("/Users/rob/Documents/Eclipse/571Search/sorted_word_doc_list.txt");
    	
    	//Index.fillDocWordFreqDfList("/Users/rob/Documents/Eclipse/571Search/doc_word_freq_df.txt");

    	//DO: sort -k1n,1 -k2,2 doc_word_freq_df.txt > sorted_doc_word_freq_df.txt
    	Index.fillDocMagTable("/Users/rob/Documents/Eclipse/571Search/sorted_doc_word_freq_df.txt", 19026);    	
    	
    	//Index.fillDocumentsTable("/Users/rob/Documents/Eclipse/571Search/html/");

    }
}





